"""
    导入标准库
"""
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression

is_print = True

"""
    读取数据
"""
dataset = pd.read_csv("../datas/Data.csv")
X = dataset.iloc[:,:-1] # 逗号前的表示行，逗号后的表示列，该式表示的是所有行，列是除最后一列不提取，即提取从0到倒数第二列
y = dataset.iloc[:,3] # 逗号前的表示行，逗号后的表示列，该式表示的是所有行，提取指定列，索引是从0开始的，即0,1,2,3,提取列3

if is_print:
    print(X)
    print("-------------------------------------")
    print(y)

"""
    空值处理
"""
X['Age'] = X['Age'].fillna(X['Age'].mean())
X['Salary'] = X['Salary'].fillna(X['Salary'].mean())

"""
    类型特征编码-独热码one_hot
"""
# 特征X
X = pd.get_dummies(X)
X = X.values
# 标签y
le = LabelEncoder()
y = le.fit_transform(y)

"""
    划分训练集、测试集
"""
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)

"""
    标准化:只X需要标准化，而标签y不需要标准化
"""
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

"""
    简单线性回归算法
"""
regressor = LinearRegression()
regressor.fit(X_train, y_train)

"""
    对测试集进行预测
"""
y_pred = regressor.predict(X_test)

"""
    计算J
"""
J = 1/X_train.shape[0] * np.sum((regressor.predict(X_train) - y_train)**2)
print("J = {}".format(J))

"""
    计算参数 w0、w1
"""
w0 = regressor.intercept_
w1 = regressor.coef_[0]
print("w0 = {}, w1 = {}".format(w0, w1))

"""
    可视化训练集拟合结果
"""
plt.figure(1)
plt.scatter(X_train, y_train, color = 'red')
plt.plot(X_train, regressor.predict(X_train), color = 'blue')
plt.title('population VS median_house_value (training set)')
plt.xlabel('population')
plt.ylabel('median_house_value')
plt.show()

"""
    可视化测试集拟合结果
"""
plt.figure(2)
plt.scatter(X_test, y_test, color = 'red')
plt.plot(X_train, regressor.predict(X_train), color = 'blue')
plt.title('population VS median_house_value (test set)')
plt.xlabel('population')
plt.ylabel('median_house_value')
plt.show()